#pragma once
#include "stream.h"
#include "task.h"
#include <concepts>
#include <cstddef>
#include <cstdint>
#include <iostream>
#include <istream>
#include <optional>
#include <sstream>
#include <string>
#include <sys/types.h>
#include <unordered_map>
#include <utility>
#include <variant>
#include <vector>

namespace cobt {
using BMap = std::unordered_map<std::string, class BEncodeData>;
using BList = std::vector<class BEncodeData>;
template <typename T>
concept ConvertibleToString = requires(T t) {
  { std::string(t) } -> std::convertible_to<std::string>;
};

// Define a concept to check if a type is convertible to int64_t
template <typename T>
concept ConvertibleToInt64 = requires(T t) {
  { static_cast<int64_t>(t) } -> std::convertible_to<int64_t>;
};

template <typename T>
concept Convertible = ConvertibleToString<T> || ConvertibleToInt64<T>;

template <typename T>
concept BEncodeDataType =
    Convertible<T> || std::same_as<T, BList> || std::same_as<T, BMap>;

class BEncodeData {
public:
  BEncodeData() = default;
  template <BEncodeDataType T>
  BEncodeData(T &&val) : m_data(std::forward<T>(val)) {}
  std::string get_raw() const { return m_raw; }

  std::string dump() const {
    std::stringstream ret;
    switch (m_data.index()) {
    case 0: {
      std::string data = std::get<0>(m_data);
      ret << data.size() << ':' << std::move(data);
      return ret.str();
    }
    case 1: {
      ret << 'i' << std::get<1>(m_data) << 'e';
      return ret.str();
    }
    case 2: {
      ret << 'l';
      for (auto &item : std::get<2>(m_data)) {
        ret << item.dump();
      }
      ret << 'e';
      return ret.str();
    }
    case 3: {
      ret << 'd';
      for (auto &[key, value] : std::get<3>(m_data)) {
        ret << key.size() << ':' << key << value.dump();
      }
      ret << 'e';
      return ret.str();
    }
    }
    return "";
  }

  BEncodeData *operator[](const std::string &key) {
    if (m_data.index() != 3) {
      return nullptr;
    }
    BMap &bm = std::get<3>(m_data);
    auto it = bm.find(key);
    if (it == bm.end()) {
      return nullptr;
    }
    return &it->second;
  }
  const BEncodeData *operator[](const std::string &key) const {
    if (m_data.index() != 3) {
      return nullptr;
    }
    const BMap &bm = std::get<3>(m_data);
    auto it = bm.find(key);
    if (it == bm.end()) {
      return nullptr;
    }
    return &it->second;
  }
  BEncodeData *operator[](size_t index) {
    if (m_data.index() != 2) {
      return nullptr;
    }
    BList &bl = std::get<2>(m_data);
    if (index >= bl.size()) {
      return nullptr;
    }
    return &bl[index];
  }
  const BEncodeData *operator[](size_t index) const {
    if (m_data.index() != 2) {
      return nullptr;
    }
    const BList &bl = std::get<2>(m_data);
    if (index >= bl.size()) {
      return nullptr;
    }
    return &bl[index];
  }
  template <BEncodeDataType T> T *as() {
    if (!std::holds_alternative<T>(m_data)) {
      return nullptr;
    }
    return &std::get<T>(m_data);
  }
  template <BEncodeDataType T> const T *as() const {
    if (!std::holds_alternative<T>(m_data)) {
      return nullptr;
    }
    return &std::get<T>(m_data);
  }

  static std::optional<BEncodeData> parse(const std::string &str) {
    std::stringstream ss(str);
    return parse(ss);
  }
  static std::optional<BEncodeData> parse(std::istream &ss) {
    char ch = ss.peek();
    if (ch >= '0' && ch <= '9') {
      return parse_string(ss);
    }
    if (ch == 'l') {
      return parse_list(ss);
    }
    if (ch == 'd') {
      return parse_dict(ss);
    }
    if (ch == 'i') {
      return parse_int(ss);
    }
    return std::nullopt;
  }
  template <class T>
  static Task<std::optional<BEncodeData>> ReadFrom(IOStream<T> &fs) {
    auto ch = co_await fs.peek();
    if (!ch) {
      co_return std::nullopt;
    }
    if (*ch >= '0' && *ch <= '9') {
      co_return co_await ReadToString(fs);
    }
    if (*ch == 'l') {
      co_return co_await ReadToList(fs);
    }
    if (*ch == 'd') {
      co_return co_await ReadToDict(fs);
    }
    if (*ch == 'i') {
      co_return co_await ReadToInt(fs);
    }
    co_return std::nullopt;
  }

private:
  std::variant<std::string, int64_t, BList, BMap> m_data;
  std::string m_raw;
  static std::optional<BEncodeData> parse_dict(std::istream &ss) {
    if (ss.peek() != 'd') {
      return std::nullopt;
    }

    BEncodeData data(std::unordered_map<std::string, BEncodeData>{});
    data.m_raw.push_back(ss.get());
    // std::cout << "{";
    while (true) {
      if (ss.peek() == 'e') {
        data.m_raw.push_back(ss.get());
        break;
      }
      auto key = parse_string(ss);
      if (!key) {
        return std::nullopt;
      }
      // std::cout << ": ";
      data.m_raw.append(key->m_raw);
      auto value = parse(ss);
      if (!value) {
        return std::nullopt;
      }
      data.m_raw.append(value->m_raw);
      std::get<3>(data.m_data)[std::get<0>(key->m_data)] =
          std::move(value.value());
      // std::cout << ", ";
    }
    // std::cout << "}";
    return data;
  }
  static std::optional<BEncodeData> parse_list(std::istream &ss) {
    if (ss.peek() != 'l') {
      return std::nullopt;
    }

    BEncodeData data(std::vector<BEncodeData>(0));
    data.m_raw.push_back(ss.get());
    // std::cout << '[';
    while (true) {
      if (ss.peek() == 'e') {
        data.m_raw.push_back(ss.get());
        break;
      }
      auto ret = parse(ss);
      if (!ret) {
        return std::nullopt;
      }
      data.m_raw.append(ret.value().m_raw);
      std::get<2>(data.m_data).emplace_back(std::move(ret.value()));
      // std::cout << ",";
    }
    // std::cout << ']';
    return data;
  }
  static std::optional<BEncodeData> parse_string(std::istream &ss) {
    std::stringstream raw;
    std::string raw_str;
    int len;
    if (!(ss >> len)) {
      return std::nullopt;
    }
    raw << len;
    if (ss.peek() != ':') {
      return std::nullopt;
    }
    raw << (char)ss.get();
    raw_str = raw.str();
    std::string ret(len, 0);
    for (int i = 0; i < len; i++) {
      ret[i] = ss.get();
      raw_str.push_back(ret[i]);
    }
    // std::cout << ret;
    BEncodeData data(std::move(ret));

    data.m_raw = std::move(raw_str);
    return data;
  }
  static std::optional<BEncodeData> parse_int(std::istream &ss) {
    std::stringstream raw;
    if (ss.peek() != 'i') {
      return std::nullopt;
    }
    raw << (char)ss.get();
    int64_t ret;
    ss >> ret;
    raw << ret;
    if (ss.peek() != 'e') {
      return std::nullopt;
    }
    raw << (char)ss.get();
    // std::cout << ret;
    BEncodeData data(std::move(ret));
    data.m_raw = raw.str();
    return data;
  }

  template <class T>
  static Task<std::optional<BEncodeData>> ReadToList(IOStream<T> &fs) {
    auto ch = co_await fs.getchar();
    if (!ch || ch.value() != 'l') {
      co_return std::nullopt;
    }
    BList res;
    while (true) {
      auto ch = co_await fs.peek();
      if (!ch) {
        co_return std::nullopt;
      }
      if (ch.value() == 'e') {
        co_await fs.getchar();
        co_return BEncodeData(std::move(res));
      }
      if (auto data = co_await ReadFrom(fs)) {
        res.push_back(std::move(*data));
      }
    }
  }
  template <class T>
  static Task<std::optional<BEncodeData>> ReadToDict(IOStream<T> &fs) {
    auto ch = co_await fs.getchar();
    if (!ch || ch.value() != 'd') {
      co_return std::nullopt;
    }
    BMap res;
    while (true) {
      auto ch = co_await fs.peek();
      if (!ch) {
        co_return std::nullopt;
      }
      if (ch.value() == 'e') {
        co_await fs.getchar();
        co_return BEncodeData(std::move(res));
      }
      auto key = co_await ReadToString(fs);
      if (!key) {
        co_return std::nullopt;
      }
      auto val = co_await ReadFrom(fs);
      if (!val) {
        co_return std::nullopt;
      }
      res[std::move(key.value())] = std::move(val.value());
    }
  }

  template <class T>
  static Task<std::optional<std::string>> ReadToString(IOStream<T> &fs) {
    auto ch = co_await fs.peek();
    if (!ch || ch.value() < '0' || ch.value() > '9') {
      co_return std::nullopt;
    }
    std::string str;
    while (true) {
      ch = co_await fs.getchar();
      if (!ch) {
        co_return std::nullopt;
      }
      if (ch.value() >= '0' && ch.value() <= '9') {
        str += ch.value();
        continue;
      }
      if (ch.value() == ':') {
        break;
      }
      co_return std::nullopt;
    }
    size_t len = std::stoi(str);
    str = co_await fs.getn(len);
    if (str.size() != len) {
      co_return std::nullopt;
    }
    co_return str;
  }
  template <class T>
  static Task<std::optional<int64_t>> ReadToInt(IOStream<T> &fs) {
    auto ch = co_await fs.getchar();
    if (!ch || ch.value() != 'i') {
      co_return std::nullopt;
    }
    std::string str;
    bool minus = false;
    while (true) {
      ch = co_await fs.getchar();
      if (!ch) {
        co_return std::nullopt;
      }
      if (ch.value() == 'e') {
        if (str.empty()) {
          co_return std::nullopt;
        }
        co_return std::stoll(str);
      }
      if (ch.value() == '-') {
        if (minus) {
          co_return std::nullopt;
        }
        str.push_back(ch.value());
        minus = true;
        continue;
      }
      if (ch.value() >= '0' && ch.value() <= '9') {
        str.push_back(ch.value());
        continue;
      }
      co_return std::nullopt;
    }
  }
};
} // namespace cobt
